# -*- coding: utf-8 -*-
"""
Created on Mon Jul 11 14:30:14 2022

@author: 潘宇浩
"""


class findEA(object):
    import numpy as np
    def __init__(self,runcal,x0,x_lb=[0,0],x_ub=[90,90],T0=1,alfa=0.8,L=5,maxgen=10,start=[]):
        import numpy as np
        self._x0=np.array(x0)*np.pi/180
        self.x_lb=np.array(x_lb)*np.pi/180
        self.x_ub=np.array(x_ub)*np.pi/180
        self._T0=T0
        self.alfa=alfa
        self.L=L
        self.maxgen=maxgen
        self.runcal=runcal
        new_start=np.array(start)*np.pi/180
        self.start=new_start.tolist()


    def Subwork(self,x,first=False):
        import os
        import numpy as np
        fb=first
        self.caldir=str(round(x[0]*180/np.pi,3))+"-"+str(round(x[1]*180/np.pi,3))
        chmk=self.makedir(x)
        if chmk == "Finished":
            os.chdir(self.caldir)
            os.system(self.runcal)           
            y=self.reader()
            os.chdir("../")
            if y == "Failed":
                return "Failed"
            else:
                self.record(x,y,first=fb)
                return y
        elif chmk == "repeat":
            os.chdir(self.caldir)
            y=self.reader()
            os.chdir("../")
            self.record(x,y,first=fb)
            return y
        else:
            return "Failed"
    
    def makedir(self,x):
        import os
        import shutil
        if self.check_repeat():
            return "repeat"
        else:
            try:
                os.mkdir(self.caldir)
                if os.path.exists('vdw_kernel.bindat'):
                    os.chdir(self.caldir)
                    os.symlink('../KPOINTS','./KPOINTS')
                    os.symlink('../POTCAR','./POTCAR')
                    os.symlink('../POSCAR','./POSCAR')
                    shutil.copy('../INCAR','./INCAR')
                    os.symlink('../vdw_kernel.bindat','./vdw_kernel.bindat')
                    os.symlink('../CHGCAR','./CHGCAR')
                else:
                    os.chdir(self.caldir)
                    os.symlink('../KPOINTS','./KPOINTS')
                    os.symlink('../POTCAR','./POTCAR')
                    os.symlink('../POSCAR','./POSCAR')    
                    shutil.copy('../INCAR','./INCAR') 
                    os.symlink('../CHGCAR','./CHGCAR')
                self.write_incar(x)
                os.chdir("..")
                return "Finished"
            except:
                return "Failed"

    def check_repeat(self):
        import os
        dirnames=os.listdir('./')
        if self.caldir in dirnames:
            return True
        else:
            return False

    def write_incar(self,x):
        import numpy as np
        fi=open("INCAR","a+")
        a=np.sin(x[0])*np.cos(x[1])
        b=np.sin(x[0])*np.sin(x[1])
        c=np.cos(x[0])
        fi.write("\nSAXIS = %.8f %.8f %.8f"%(a,b,c))
        fi.close()
        return

    def reader(self):
        if self.check_conv():
            sout_f=open('OSZICAR','rt')
            sout_lines=sout_f.readlines()
            y=float(sout_lines[-1].split()[4])
            return y
        else:
            return "Failed"

    def check_conv(self):
        try:
            sout_f=open('OSZICAR','rt')
            sout_num=int(sout_f.readlines()[-2].split()[1])
            sout_f.close()
            nelm=0
            incar_f=open('INCAR','rt')
            incar_lines=incar_f.readlines()
            for ss in incar_lines:
                if 'NELM' in ss and "NELMDL" not in ss and "NELMIN" not in ss and "#" not in ss and "!" not in ss:
                    for n in ss.split():
                        try:
                            nelm=int(n)
                        except:
                            continue
            if nelm == 0:
                nelm=60
            incar_f.close()
            if sout_num < nelm:
                return True
            else:
                return False
        except:
            return False
    
    
    def record(self,x,y,first=False,mid=False):
        import numpy as np
        if mid:
            fr=open("record.txt","a+")
            fr.write("\t(%.3f\t%.3f)"%(x[0]*180/np.pi,x[1]*180/np.pi))
            fr.close()
        elif first:
            import os
            if os.path.exists("record.txt"):
                os.remove("record.txt")
            fr=open("record.txt","a+")
            fr.write("\n%.3f\t%.3f\t%.4f\t%.4f\t%.8f\t%.8f\t%.8f\t%.8f"%(x[0]*180/np.pi,x[1]*180/np.pi,x[0],x[1],np.sin(x[0])*np.cos(x[1]),np.sin(x[
0])*np.sin(x[1]),np.cos(x[0]),y))
            fr.close()
        else:
            fr=open("record.txt","a+")
            fr.write("\n%.3f\t%.3f\t%.4f\t%.4f\t%.8f\t%.8f\t%.8f\t%.8f"%(x[0]*180/np.pi,x[1]*180/np.pi,x[0],x[1],np.sin(x[0])*np.cos(x[1]),np.sin(x[
0])*np.sin(x[1]),np.cos(x[0]),y))
            fr.close()
        return
    
    def run(self):
        import numpy as np
        x=self._x0
        y=self.Subwork(x,first=True)
        if y == "Failed":
            fr=open("record.txt","a+")
            fr.write("Failed to Make Folder OR Calculation was Not Convergent. Please Check!")
            fr.close()
            return
        best_x=x
        best_y=y
        self.best_x=best_x.copy()
        self.best_y=best_y
        self.record(best_x,best_y,mid=True)
        startup=[[0,np.pi/2],[np.pi/2,0],[np.pi/2,np.pi/2]]+self.start
        for start_x in startup:
            y_startup=self.Subwork(start_x)
            if y == "Failed":
                fr=open("record.txt","a+")
                fr.write("Failed to Make Folder OR Calculation was Not Convergent. Please Check!")
                fr.close()
                return
            if y_startup<best_y:
                x=start_x.copy()
                y=y_startup
                best_x=start_x.copy()
                best_y=y_startup
                self.best_x=best_x.copy()
                self.best_y=best_y
                self.record(best_x,best_y,mid=True)
        T=self._T0
      
        for it in range(self.maxgen):
            for i in range(self.L):
                r=np.random.randn(np.size(x))
                z=r/np.sqrt(np.sum(r**2))
#                x_new=np.round(x+z*T,8)
                x_new=x+z*T
                for j in range(np.size(x)):
                    if x_new[j]<self.x_lb[j]:
                        r2=np.random.rand()
#                        x_new[j]=np.round(r2*x[j]+(1-r2)*self.x_lb[j],8)
                        x_new[j]=r2*x[j]+(1-r2)*self.x_lb[j]
                    elif x_new[j]>self.x_ub[j]:
                        r2=np.random.rand()
#                        x_new[j]=np.round(r2*x[j]+(1-r2)*self.x_ub[j],8)
                        x_new[j]=r2*x[j]+(1-r2)*self.x_ub[j]
                y_new=self.Subwork(x_new)
                
                if y_new == "Failed":
                    fr=open("record.txt","a+")
                    fr.write("Failed to Make Folder OR Calculation was Not Convergent. Please Check!")
                    fr.close()
                    return
                
                if y_new<best_y:
                    x=x_new.copy()
                    y=y_new
                    best_x=x_new.copy()
                    best_y=y_new
                    self.best_x=best_x.copy()
                    self.best_y=best_y
                    self.record(best_x,best_y,mid=True)
                else:
                    p=np.exp(-(y_new-y)*100/T)
                    if np.random.rand() < p:
                        x=x_new.copy()
                        y=y_new
            T *= self.alfa
        self.best_x=best_x.copy()
        self.best_y=best_y
        fr=open("record.txt","a+")
        fr.write("\nBest Result:")
        fr.close()
        self.record(self.best_x, self.best_y)
        return

if __name__=="__main__":
    import sys
    import numpy as np
    input_f=open('input','rt')
    input_lines=input_f.readlines()
#    rrun=input_lines[0][:-1]
    rrun="mpirun -f ../host.txt -n "+sys.argv[1]+" "+input_lines[0][:-1]+" | tee sout"
    x0=[float(input_lines[1].split()[0]),float(input_lines[2].split()[0])]
    new_start=[[float(input_lines[1].split()[i]),float(input_lines[2].split()[i])] for i in range(1,len(input_lines[1].split()))]
    x_lb=[float(input_lines[3]),float(input_lines[4])]
    x_ub=[float(input_lines[5]),float(input_lines[6])]
    T0=float(input_lines[7])
    alfa=float(input_lines[8])
    L=int(input_lines[9])
    maxgen=int(input_lines[10])
    ea=findEA(rrun,x0,x_lb,x_ub,T0,alfa,L,maxgen,new_start)
    ea.run()